package com.jhm.ai.util;

import com.jhm.ai.domain.ChatRequest;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class SseEmitterUtil {
    private static final Logger log = LoggerFactory.getLogger(SseEmitterUtil.class);
    private static final Map<String, SseEmitter> EMITTER_MAP = new ConcurrentHashMap<>();
    private static final Map<String, ChatRequest> REQUEST_MAP = new ConcurrentHashMap<>();
    private static final long TIMEOUT = 3600000L; // 1小时超时

    public static void storeRequest(String clientId, ChatRequest request) {
        REQUEST_MAP.put(clientId, request);
    }

    public static ChatRequest getRequest(String clientId) {
        return REQUEST_MAP.remove(clientId); // 获取后立即删除
    }

    public static SseEmitter createEmitter(String clientId) {
        SseEmitter emitter = new SseEmitter(TIMEOUT);
        EMITTER_MAP.put(clientId, emitter);

        emitter.onCompletion(() -> {
            log.info("SSE connection completed for client: {}", clientId);
            EMITTER_MAP.remove(clientId);
            REQUEST_MAP.remove(clientId);
        });

        emitter.onTimeout(() -> {
            log.info("SSE connection timeout for client: {}", clientId);
            emitter.complete();
            EMITTER_MAP.remove(clientId);
            REQUEST_MAP.remove(clientId);
        });

        emitter.onError(ex -> {
            log.error("SSE connection error for client: {}", clientId, ex);
            emitter.complete();
            EMITTER_MAP.remove(clientId);
            REQUEST_MAP.remove(clientId);
        });

        return emitter;
    }

    public static void removeEmitter(String clientId) {
        SseEmitter emitter = EMITTER_MAP.get(clientId);
        if (emitter != null) {
            emitter.complete();
            EMITTER_MAP.remove(clientId);
        }
    }

    public static void send(String clientId, Object data) {
        SseEmitter emitter = EMITTER_MAP.get(clientId);
        if (emitter != null) {
            try {
                emitter.send(SseEmitter.event()
                        .data(data)
                        .build());
            } catch (IOException e) {
                log.error("Error sending SSE event to client: {}", clientId, e);
                emitter.completeWithError(e);
                EMITTER_MAP.remove(clientId);
            }
        }
    }

    public static boolean hasEmitter(String clientId) {
        return EMITTER_MAP.containsKey(clientId);
    }

    public static void complete(String clientId) {
        SseEmitter emitter = EMITTER_MAP.get(clientId);
        if (emitter != null) {
            emitter.complete();
            EMITTER_MAP.remove(clientId);
        }
    }
} 